{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4bdd4942",
   "metadata": {},
   "outputs": [],
   "source": [
    "# step1 导入相关包\n",
    "import copy\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import time\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21a4f75d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 设置随机种子 确保我们得到的结果是相同的, 使得这个过程变得确定性,消除任何随机性的影响。\n",
    "torch.manual_seed(42)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3fb6a61",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 定义一个测试模型, 最简单的测试模型，仅有3层， 嵌入层、线性层、LLM 头层, 也不做任何的注意力\n",
    "class TestModel(torch.nn.Module):\n",
    "    def __init__(self.hidden_size):\n",
    "        super().__init__()\n",
    "        self.embedding = torch.nn.Embedding(10, hidden_size)\n",
    "        self.linear = torch.nn.Linear(hidden_size, hidden_size)\n",
    "        self.lm_head = torch.nn.Linear(hidden_size, 10)\n",
    "        \n",
    "    \n",
    "    def forward(self, input_ids):\n",
    "        x = self.embedding(input_ids)\n",
    "        x = self.linear(x)\n",
    "        x = self.lm_head(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3be109d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 设置一个很大的 hidden_size\n",
    "hidden_size = 1024\n",
    "model = TestModel(hidden_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f117bcb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 设置一个虚拟的 input_ids\n",
    "input_ids = torch.LongTensor([[0, 1, 2 , 3, 4, 5, 6, 7]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa625c59",
   "metadata": {},
   "outputs": [],
   "source": [
    "# detokenizer example\n",
    "detokenizer = [\n",
    "    \"red\",\n",
    "    \"orange\",\n",
    "    \"yellow\",\n",
    "    \"green\",\n",
    "    \"blue\",\n",
    "    \"indigo\",\n",
    "    \"violet\",\n",
    "    \"magenta\",\n",
    "    \"marigold\",\n",
    "    \"chartreuse\",\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38a1913b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_token(model, **kwargs):\n",
    "    with torch.no_grad():\n",
    "        outputs = model(**kwargs)\n",
    "    \n",
    "    logtis = outputs.logits\n",
    "    last_logits = logits[:, -1, :]\n",
    "    next_token_ids = last_logits.argmax(dim=1)\n",
    "    return [detokenizer[token_id] for token_id in next_token_ids]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "443baecb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 先假设一个input tensor\n",
    "x = torch.randn(1, 8, 1024)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ff93b04",
   "metadata": {},
   "outputs": [],
   "source": [
    "lora_a = torch.randn(1024, 2)\n",
    "lora_b = torch.randn(2, 1024)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33ca249a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 模型原始权重\n",
    "W = model.linear.weight"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0a64bb6",
   "metadata": {},
   "outputs": [],
   "source": [
    "W2 = lora_a @ lora_b\n",
    "\n",
    "W2.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f23109b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compare number of elements of A and B with number of elements of W\n",
    "# W here has shape (hidden_size, hidden_size)\n",
    "# LoRA 配置的时候，可以设置 秩： r, 若增加r 的值，参数的百分比应该接近 W， 若隐藏大小 增加，这个值会降低\n",
    "tora_numel = tora_a.numel() + tora_b.numel()\n",
    "base_numel = W.numel()\n",
    "print(\"|A+B| / |W|:\", tora_numel / base_numel)\n",
    "\n",
    "# |A+B| / |W|: 0.00390625\n",
    "# 它不到 W中 参数量的 百分之一"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5cf47037",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 接下来计算lora 层的输出\n",
    "base_output = model.linear(X)\n",
    "\n",
    "# 计算lora 层的输出\n",
    "lora_output = X @ lora_a @ lora_b\n",
    "\n",
    "# sum\n",
    "total_output = base_output + lora_output\n",
    "\n",
    "total_output.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d209e01",
   "metadata": {},
   "outputs": [],
   "source": [
    "# dummy input tensor\n",
    "# shape: (batch_size, sequence_length, hidden_size)\n",
    "X = torch.randn(1, 8, 1024)\n",
    "\n",
    "# LoRA A and B tensors\n",
    "# A has shape (hidden_size, rank)\n",
    "# B has shape (rank, hidden_size)\n",
    "lora_a = torch.randn(1024, 2)\n",
    "lora_b = torch.randn(2, 1024)\n",
    "\n",
    "W = model.linear.weight\n",
    "\n",
    "W.shape\n",
    "\n",
    "torch.Size([1024, 1024])\n",
    "\n",
    "W2 = lora_a @ lora_b\n",
    "\n",
    "W2.shape\n",
    "\n",
    "torch.Size([1024, 1024])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5c3f4f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 定义lora 层， 11:32 分\n",
    "# 它接收一个模型中的原始线性层， 并接受一个 r, 将使用这些输入来存储基础层，并创建 lora_A  和  lora_B\n",
    "class Loralayer(torch.nn.Module):\n",
    "    def __init__(self, base_layer, r):\n",
    "        super().__init__()\n",
    "        self.base_layer = base_layer\n",
    "        d_in, d_out = self.base_layer.weight.shape\n",
    "        self.lora_a = torch.randn(d_in, r)\n",
    "        self.lora_b = torch.randn(r, d_out)\n",
    "\n",
    "    def forward(self, x):\n",
    "        y1 = self.base_layer(x)\n",
    "        y2 = x @ self.lora_a @ self.lora_b\n",
    "        return y1 + y2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34bea8a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 我们用新的lora 层替换了模型的 线性层，\n",
    "# wrap the linear layer of our toy model, use rank 2\n",
    "lora_layer = Loralayer(model.linear, 2)\n",
    "lora_layer(X).shape\n",
    "\n",
    "torch.Size([1, 8, 1024])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de56e5e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# generate one token\n",
    "next_token = generate_token(model, input_ids=input_ids)[0]\n",
    "next_token"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "tuning_env",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.10.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
